package rerac.feature;

import java.util.ArrayList;
import java.util.List;

import rerac.protos.Corpus.Document;
import rerac.protos.Corpus.Document.Annotation;
import rerac.protos.Corpus.Document.AnnotationType;
import rerac.protos.Corpus.Document.Method;
import rerac.protos.Corpus.Document.Token;

import edu.stanford.nlp.ling.HasWord;
import edu.stanford.nlp.parser.lexparser.LexicalizedParser;
import edu.stanford.nlp.trees.GrammaticalStructure;
import edu.stanford.nlp.trees.GrammaticalStructureFactory;
import edu.stanford.nlp.trees.PennTreebankLanguagePack;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.trees.TreebankLanguagePack;
import edu.stanford.nlp.trees.TypedDependency;

/**
 * 
 * This is a wrapper class for using the Stanford parser.
 * 
 * @author Benjamin Roth
 *
 */
public class DocumentParser {
  
  /**
   * Wrapper for information about words.
   * 
   * @author Benjamin Roth
   *
   */
  class Word implements HasWord{
    private static final long serialVersionUID = 1L;
    String word;    
    Word(String w) {
      setWord(w);
    }
    
    @Override
    public void setWord(String w) {
      word = w;
    }
    @Override
    public String word() {
      return word;
    }    
  }
  
  // Parser used.
  LexicalizedParser lp;
  // Used to create the dependency strucure.
  GrammaticalStructureFactory gsf;
  
  /**
   * This creates a parser with a given grammar and standard option settings.
   * 
   * @param grammar the file name of the grammar model.
   */
  public DocumentParser(String grammar) {
    this(grammar, 
        new String[] { "-maxLength", "80", "-retainTmpSubcategories",
            "-outputFormat", "typedDependenciesCollapsed" });
  }
  
  /**
   * 
   * This creates a parser with a given grammar and given option settings.
   * 
   * @param grammar the file name of the grammar model.
   * @param options the array of option strings.
   */
  DocumentParser(String model, String[] options) {
    lp = new LexicalizedParser(model);
    lp.setOptionFlags(options);
    TreebankLanguagePack tlp = new PennTreebankLanguagePack();
    gsf = tlp.grammaticalStructureFactory();
  }
  
  /**
   * This takes a document parses it and creates a new document with added 
   * parse information which is returned.
   * 
   * @param doc the document to be parsed.
   * @return the document with added parse information.
   */
  public Document parse(Document doc) {
    Document.Builder docParse = Document.newBuilder(doc);

    int parserIdx = docParse.getMethodCount();
    String parserName = this.getClass().getName();
    AnnotationType parserType = AnnotationType.SYNTAX_HEAD;
    
    Method.Builder m = 
      Method.newBuilder().setId(parserName).setType(parserType);
    // In general, the parser can add 0, 1 or more annotations to a token. 
    docParse.setConsistentMethod(false);

    docParse.addMethod(m);
    
    List<Word> wordList = new ArrayList<Word>();
    for (Token tok : doc.getTokenList()) {
      wordList.add(new Word(tok.getText()));
    }
    Tree parse = null;
    if (lp.parse(wordList)) {
      parse = lp.getBestParse();
    }

    GrammaticalStructure gs = gsf.newGrammaticalStructure(parse);
    for(TypedDependency tdp : gs.typedDependenciesCollapsed()) {
      int depIdx = tdp.dep().index() - 1;
      int govIdx = tdp.gov().index() - 1;
      Token.Builder tok = Token.newBuilder(docParse.getToken(depIdx));
      Annotation.Builder anno = Annotation.newBuilder();
      anno.setType(parserType);
      anno.setTargetToken(govIdx);
      anno.setMethodIndex(parserIdx);
      anno.setText(tdp.reln().getShortName());
      tok.addAnnotation(anno);
      
      docParse.setToken(depIdx, tok);
    }
    return docParse.build();
  }

}
